package com.test.springai.service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

/**
 * Created with IntelliJ IDEA.
 *
 * @description: 默认的查询翻译器实现。 负责利用 LLM 将用户查询从源语言翻译到目标语言。
 * @author: liuziyang
 * @since: 9/3/25 2:41 PM
 * @modifiedBy:
 * @version: 1.0
 */
@Slf4j
@Component
public class DefaultQueryTranslator implements TranslationQueryTransformer {
  private final ChatClient chatClient;

  /** 用于异步 LLM 调用的线程池 */
  private final Executor threadPoolTaskExecutor;

  /** 提示模板：用于将查询从源语言翻译到目标语言 {source_language} 和 {target_language} 是占位符，可以在运行时动态填充 */
  private final String TRANSLATION_PROMPT =
      """
            你是一个专业的语言翻译器。请将以下文本从 {source_language} 翻译成 {target_language}。
            只返回翻译后的文本，不要添加任何额外信息、解释或前缀。
            文本: {text_to_translate}
            """;

  /**
   * 构造函数，用于初始化 {@code DefaultQueryTranslator}。
   *
   * @param chatClientBuilder 用于构建 {@code ChatClient} 实例的 {@code ChatClient.Builder}。
   * @param threadPoolTaskExecutor 用于异步任务执行的 {@code Executor} 实例（例如，一个定制的线程池）。
   */
  public DefaultQueryTranslator(
      ChatClient.Builder chatClientBuilder,
      @Qualifier("applicationTaskExecutor") Executor threadPoolTaskExecutor) {
    this.chatClient = chatClientBuilder.build();
    this.threadPoolTaskExecutor = threadPoolTaskExecutor;
  }

  /**
   * 异步翻译用户查询。 该方法会向 LLM 发送一个提示，指导其将原始查询从源语言翻译到目标语言。
   *
   * @param originalQuery 用户原始查询字符串。
   * @param sourceLanguage 源语言（例如 "中文", "English"）。
   * @param targetLanguage 目标语言（例如 "英文", "French"）。
   * @return 包含翻译后查询字符串的 {@code CompletableFuture}。
   */
  @Override
  public CompletableFuture<String> translateQuery(
      String originalQuery, String sourceLanguage, String targetLanguage) {
    // 在指定的线程池中执行异步任务
    return CompletableFuture.supplyAsync(
        () -> {
          log.info(
              "Translating query from {} to {}: {}", sourceLanguage, targetLanguage, originalQuery);
          SystemPromptTemplate translationTemplate = new SystemPromptTemplate(TRANSLATION_PROMPT);
          // 填充模板中的占位符 {source_language}, {target_language}, {text_to_translate}
          String systemMessageContent =
              translationTemplate
                  .createMessage(
                      Map.of(
                          "source_language", sourceLanguage,
                          "target_language", targetLanguage,
                          "text_to_translate", originalQuery))
                  .getText();
          Prompt prompt = new Prompt(List.of(new UserMessage(systemMessageContent)));
          String translatedQuery = chatClient.prompt(prompt).call().content();
          log.debug("Translated query raw response: {}", translatedQuery);
          // LLM 可能会返回空字符串或不符合预期的内容，这里做一下处理
          if (translatedQuery != null && !translatedQuery.trim().isEmpty()) {
            return translatedQuery.trim();
          }
          // 如果翻译失败（LLM返回空或无效），则回退到原始查询
          log.warn(
              "Translation failed for query '{}' from {} to {}, returning original query.",
              originalQuery,
              sourceLanguage,
              targetLanguage);
          return originalQuery;
        },
        threadPoolTaskExecutor);
  }
}
